import pandas as pda
import numpy as npy
import time
from sklearn.model_selection import train_test_split
from sklearn import datasets

def prepare_data(test_size, random_state):
        iris_data = datasets.load_iris()
        x = iris_data.data
        y = iris_data.target
        return train_test_split(x, y, test_size = test_size, random_state = random_state)

class Bayes:
    def __init__(self):
        self.train_v_len = -1
        self.label_p_val = dict()
        self.vector_val = dict()

    def fit(self, x_train:list, y_train:list):
        print("training begin...")
        if(len(x_train) != len(y_train)):
            raise ValueError("Length of input data is not correct!")

        self.train_v_len = len(x_train[0])
        label_len = len(y_train)
        label_set = set(y_train)

        for label in label_set:
            ratio = y_train.count(label)/label_len
            self.label_p_val[label] = ratio

        for vector, label in zip(x_train, y_train):
            if(label not in self.vector_val):
                self.vector_val[label] = []
            self.vector_val[label].append(vector)
        print("training end...")

    def test(self, x_test, label_test):
        if(self.train_v_len == -1):
            raise ValueError("Please train first!")

        label_ratio = dict()
        label_set = set(label_test)

        for label in label_set:
            p = 1
            label_p_val = self.label_p_val[label]
            train_vectors = self.vector_val[label]
            v_len = len(train_vectors)
            t_vectors = npy.array(train_vectors).T

            for i in range(0, len(x_test)):
                vector = list(t_vectors[i])
                p = p * (vector.count(x_test[i])/v_len)

            label_ratio[label] = p*label_p_val
        sort_label_ratio = sorted(label_ratio, key = lambda x:label_ratio[x], reverse = True)
        return sort_label_ratio[0]

def run_bayes():
    x_train,x_test,y_train,y_test = prepare_data(0.1, 20)
    begin_time = time.time()

    bys = Bayes()
    bys.fit(x_train.tolist(),y_train.tolist())

    test_predict = []
    x_test_list = x_test.tolist()
    for i in range(0, len(x_test_list)):
        predict = bys.test(x_test_list[i], y_test.tolist())
        test_predict.append(predict)

    true = test_predict == y_test    
    print("Correct Ratio %: " + str((npy.sum(true)/len(true)) * 100) + "%")
    end_time =time.time()
    print("Elapsed Time: " + str((end_time - begin_time) * 10009) + " milliseconds")

run_bayes()
